# -*- coding:utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F


class CNNMatcher(nn.Module):
    def __init__(self, emb_vectors, kernels, num_filters, dropout_p=0.0):
        super(CNNMatcher, self).__init__()

        self.vocab_size = emb_vectors.size(0)
        self.emb_size = emb_vectors.size(1)

        self.embeddings = nn.Embedding.from_pretrained(emb_vectors)

        self.conv_layers = nn.ModuleList([])
        for kernel in kernels:
            conv_layer_i = nn.Sequential(
                nn.Conv2d(1, num_filters, kernel_size=(kernel, self.emb_size)),
                nn.ReLU()
            )
            self.conv_layers.add_module(name='conv-{}'.format(kernel),
                                        module=conv_layer_i)
        self.dropout = nn.Dropout(dropout_p)
        pass

    def conv_and_max_pooling(self, inputs):
        conv_outputs = []
        for conv_layer in self.conv_layers:
            output_i = conv_layer(inputs.unsqueeze(1)).squeeze(3)
            output_i = F.max_pool1d(output_i, output_i.size(2)).squeeze(2)
            # print(output_i.size())
            conv_outputs += [output_i]
        return torch.cat(conv_outputs, dim=1)

    def forward(self, query_inputs, candidate_inputs, **kwargs):
        # get embedded
        query_emb_inputs = self.embeddings(query_inputs)
        query_emb_inputs = self.dropout(query_emb_inputs)

        candidate_emb_inputs = self.embeddings(candidate_inputs)
        candidate_emb_inputs = self.dropout(candidate_emb_inputs)

        query_conv_output = self.conv_and_max_pooling(query_emb_inputs)
        candidate_conv_output = self.conv_and_max_pooling(candidate_emb_inputs)

        similarity = F.cosine_similarity(query_conv_output, candidate_conv_output, dim=1)
        return similarity


if __name__ == '__main__':
    embs = torch.randn(100, 64)
    matcher = CNNMatcher(embs, kernels=[2, 3, 4], num_filters=16, dropout_p=0.2)
    print(matcher)

    input1 = torch.LongTensor([[1, 3], [20, 30]])
    input2 = torch.LongTensor([[1, 30], [2, 3]])
    print(matcher.forward(input1, input2))
    pass
